// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/serializers/presto_serializer_estimation_utils.h>

#include <pollux/vector/flat_vector.h>
#include <pollux/vector/vector_type_utils.h>

namespace kumo::pollux::serializer::presto::detail {
    namespace {
        template<TypeKind Kind>
        void estimateFlattenedConstantSerializedSize(
            const BaseVector *vector,
            const melon::Range<const IndexRange *> &ranges,
            vector_size_t **sizes,
            Scratch &scratch) {
            POLLUX_CHECK_EQ(vector->encoding(), VectorEncoding::Simple::CONSTANT);
            using T = typename KindToFlatVector<Kind>::WrapperType;
            auto *constant_vector = vector->as<ConstantVector<T> >();
            if (constant_vector->value_vector()) {
                estimateWrapperSerializedSize(ranges, sizes, vector, scratch);
                return;
            }

            int32_t elementSize = sizeof(T);
            if (constant_vector->is_null_at(0)) {
                elementSize = 1;
            } else if constexpr (std::is_same_v<T, StringView>) {
                elementSize = constant_vector->value_at(0).size();
            }
            for (int32_t i = 0; i < ranges.size(); ++i) {
                *sizes[i] += elementSize * ranges[i].size;
            }
        }

        void estimateBiasedSerializedSize(
            const BaseVector *vector,
            const melon::Range<const IndexRange *> &ranges,
            vector_size_t **sizes) {
            auto valueSize = vector->type()->cppSizeInBytes();
            if (vector->may_have_nulls()) {
                auto rawNulls = vector->raw_nulls();
                for (int32_t i = 0; i < ranges.size(); ++i) {
                    auto end = ranges[i].begin + ranges[i].size;
                    int32_t numValues = bits::countBits(rawNulls, ranges[i].begin, end);
                    *(sizes[i]) += numValues * valueSize + bits::nbytes(ranges[i].size);
                }
            } else {
                for (int32_t i = 0; i < ranges.size(); ++i) {
                    *(sizes[i]) += ranges[i].size * valueSize;
                }
            }
        }

        template<TypeKind Kind>
        void estimateFlatSerializedSize(
            const BaseVector *vector,
            const melon::Range<const vector_size_t *> &rows,
            vector_size_t **sizes,
            Scratch &scratch) {
            const auto valueSize = vector->type()->cppSizeInBytes();
            const auto numRows = rows.size();
            if (vector->may_have_nulls()) {
                auto rawNulls = vector->raw_nulls();
                ScratchPtr<uint64_t, 4> nullsHolder(scratch);
                ScratchPtr<int32_t, 64> nonNullsHolder(scratch);
                auto nulls = nullsHolder.get(bits::nwords(numRows));
                simd::gatherBits(rawNulls, rows, nulls);
                auto nonNulls = nonNullsHolder.get(numRows);
                const auto numNonNull = simd::indicesOfSetBits(nulls, 0, numRows, nonNulls);
                for (int32_t i = 0; i < numNonNull; ++i) {
                    *sizes[nonNulls[i]] += valueSize;
                }
            } else {
                POLLUX_UNREACHABLE("Non null fixed width case handled before this");
            }
        }

        void estimateFlatSerializedSizeVarcharOrVarbinary(
            const BaseVector *vector,
            const melon::Range<const vector_size_t *> &rows,
            vector_size_t **sizes,
            Scratch &scratch) {
            const auto numRows = rows.size();
            auto strings = static_cast<const FlatVector<StringView> *>(vector);
            auto rawNulls = strings->raw_nulls();
            auto rawValues = strings->rawValues();
            if (!rawNulls) {
                for (auto i = 0; i < rows.size(); ++i) {
                    *sizes[i] += rawValues[rows[i]].size();
                }
            } else {
                ScratchPtr<uint64_t, 4> nullsHolder(scratch);
                ScratchPtr<int32_t, 64> nonNullsHolder(scratch);
                auto nulls = nullsHolder.get(bits::nwords(numRows));
                simd::gatherBits(rawNulls, rows, nulls);
                auto *nonNulls = nonNullsHolder.get(numRows);
                auto numNonNull = simd::indicesOfSetBits(nulls, 0, numRows, nonNulls);

                for (int32_t i = 0; i < numNonNull; ++i) {
                    *sizes[nonNulls[i]] += rawValues[rows[nonNulls[i]]].size();
                }
            }
        }

        template<>
        void estimateFlatSerializedSize<TypeKind::VARCHAR>(
            const BaseVector *vector,
            const melon::Range<const vector_size_t *> &rows,
            vector_size_t **sizes,
            Scratch &scratch) {
            estimateFlatSerializedSizeVarcharOrVarbinary(vector, rows, sizes, scratch);
        }

        template<>
        void estimateFlatSerializedSize<TypeKind::VARBINARY>(
            const BaseVector *vector,
            const melon::Range<const vector_size_t *> &rows,
            vector_size_t **sizes,
            Scratch &scratch) {
            estimateFlatSerializedSizeVarcharOrVarbinary(vector, rows, sizes, scratch);
        }

        template<>
        void estimateFlatSerializedSize<TypeKind::OPAQUE>(
            const BaseVector *,
            const melon::Range<const vector_size_t *> &,
            vector_size_t **,
            Scratch &) {
            POLLUX_FAIL("Opaque type support is not implemented.");
        }

        template<TypeKind Kind>
        void estimateFlattenedConstantSerializedSize(
            const BaseVector *vector,
            const melon::Range<const vector_size_t *> &rows,
            vector_size_t **sizes,
            Scratch &scratch) {
            POLLUX_CHECK_EQ(vector->encoding(), VectorEncoding::Simple::CONSTANT);

            using T = typename KindToFlatVector<Kind>::WrapperType;
            auto *constant_vector = vector->as<ConstantVector<T> >();
            int32_t elementSize = sizeof(T);
            if (constant_vector->is_null_at(0)) {
                elementSize = 1;
            } else if (vector->value_vector()) {
                const auto *values = constant_vector->wrapped_vector();
                vector_size_t *sizePtr = &elementSize;
                const vector_size_t singleRow = constant_vector->wrapped_index(0);
                estimateSerializedSizeInt(
                    values,
                    melon::Range<const vector_size_t *>(&singleRow, 1),
                    &sizePtr,
                    scratch);
            } else if constexpr (std::is_same_v<T, StringView>) {
                elementSize = constant_vector->value_at(0).size();
            }
            for (int32_t i = 0; i < rows.size(); ++i) {
                *sizes[i] += elementSize;
            }
        }

        void estimateWrapperSerializedSize(
            const melon::Range<const vector_size_t *> &rows,
            vector_size_t **sizes,
            const BaseVector *wrapper,
            Scratch &scratch) {
            ScratchPtr<vector_size_t, 1> innerRowsHolder(scratch);
            ScratchPtr<vector_size_t *, 1> innerSizesHolder(scratch);
            const int32_t numRows = rows.size();
            int32_t numInner = 0;
            auto *innerRows = innerRowsHolder.get(numRows);
            auto *innerSizes = sizes;
            const BaseVector *wrapped;
            if (wrapper->encoding() == VectorEncoding::Simple::DICTIONARY &&
                !wrapper->raw_nulls()) {
                // Dictionary with no nulls.
                auto *indices = wrapper->wrap_info()->as<vector_size_t>();
                wrapped = wrapper->value_vector().get();
                simd::transpose(indices, rows, innerRows);
                numInner = numRows;
            } else {
                wrapped = wrapper->wrapped_vector();
                innerSizes = innerSizesHolder.get(numRows);
                for (int32_t i = 0; i < rows.size(); ++i) {
                    if (!wrapper->is_null_at(rows[i])) {
                        innerRows[numInner] = wrapper->wrapped_index(rows[i]);
                        innerSizes[numInner] = sizes[i];
                        ++numInner;
                    }
                }
            }
            if (numInner == 0) {
                return;
            }

            estimateSerializedSizeInt(
                wrapped,
                melon::Range<const vector_size_t *>(innerRows, numInner),
                innerSizes,
                scratch);
        }

        void estimateBiasedSerializedSize(
            const BaseVector *vector,
            const melon::Range<const vector_size_t *> &rows,
            vector_size_t **sizes,
            Scratch &scratch) {
            POLLUX_UNSUPPORTED();
        }
    } // namespace

    void estimateFlatSerializedSizeVarcharOrVarbinary(
        const BaseVector *vector,
        const melon::Range<const IndexRange *> &ranges,
        vector_size_t **sizes) {
        auto strings = static_cast<const FlatVector<StringView> *>(vector);
        auto rawNulls = strings->raw_nulls();
        auto rawValues = strings->rawValues();
        for (int32_t i = 0; i < ranges.size(); ++i) {
            auto end = ranges[i].begin + ranges[i].size;
            int32_t numNulls = 0;
            int32_t bytes = 0;
            for (int32_t offset = ranges[i].begin; offset < end; ++offset) {
                if (rawNulls && bits::isBitNull(rawNulls, offset)) {
                    ++numNulls;
                } else {
                    bytes += sizeof(int32_t) + rawValues[offset].size();
                }
            }
            *(sizes[i]) += bytes + bits::nbytes(numNulls) + 4 * numNulls;
        }
    }

    void estimateSerializedSizeInt(
        const BaseVector *vector,
        const melon::Range<const IndexRange *> &ranges,
        vector_size_t **sizes,
        Scratch &scratch) {
        switch (vector->encoding()) {
            case VectorEncoding::Simple::FLAT:
                POLLUX_DYNAMIC_SCALAR_TYPE_DISPATCH_ALL(
                    estimateFlatSerializedSize,
                    vector->type_kind(),
                    vector,
                    ranges,
                    sizes);
                break;
            case VectorEncoding::Simple::CONSTANT:
                POLLUX_DYNAMIC_TYPE_DISPATCH_ALL(
                    estimateFlattenedConstantSerializedSize,
                    vector->type_kind(),
                    vector,
                    ranges,
                    sizes,
                    scratch);
                break;
            case VectorEncoding::Simple::DICTIONARY:
            case VectorEncoding::Simple::SEQUENCE:
                estimateWrapperSerializedSize(ranges, sizes, vector, scratch);
                break;
            case VectorEncoding::Simple::BIASED:
                estimateBiasedSerializedSize(vector, ranges, sizes);
                break;
            case VectorEncoding::Simple::ROW: {
                std::vector<IndexRange> childRanges;
                std::vector<vector_size_t *> childSizes;
                for (int32_t i = 0; i < ranges.size(); ++i) {
                    auto begin = ranges[i].begin;
                    auto end = begin + ranges[i].size;
                    for (auto offset = begin; offset < end; ++offset) {
                        *sizes[i] += sizeof(int32_t);
                        if (!vector->is_null_at(offset)) {
                            childRanges.push_back(IndexRange{offset, 1});
                            childSizes.push_back(sizes[i]);
                        }
                    }
                }
                auto row_vector = vector->as<RowVector>();
                auto &children = row_vector->children();
                for (auto &child: children) {
                    if (child) {
                        estimateSerializedSizeInt(
                            child.get(),
                            melon::Range(childRanges.data(), childRanges.size()),
                            childSizes.data(),
                            scratch);
                    }
                }
                break;
            }
            case VectorEncoding::Simple::MAP: {
                auto map_vector = vector->as<MapVector>();
                std::vector<IndexRange> childRanges;
                std::vector<vector_size_t *> childSizes;
                expandRepeatedRanges(
                    map_vector,
                    map_vector->rawOffsets(),
                    map_vector->rawSizes(),
                    ranges,
                    sizes,
                    &childRanges,
                    &childSizes);
                estimateSerializedSizeInt(
                    map_vector->mapKeys().get(), childRanges, childSizes.data(), scratch);
                estimateSerializedSizeInt(
                    map_vector->mapValues().get(),
                    childRanges,
                    childSizes.data(),
                    scratch);
                break;
            }
            case VectorEncoding::Simple::ARRAY: {
                auto array_vector = vector->as<ArrayVector>();
                std::vector<IndexRange> childRanges;
                std::vector<vector_size_t *> childSizes;
                expandRepeatedRanges(
                    array_vector,
                    array_vector->rawOffsets(),
                    array_vector->rawSizes(),
                    ranges,
                    sizes,
                    &childRanges,
                    &childSizes);
                estimateSerializedSizeInt(
                    array_vector->elements().get(),
                    childRanges,
                    childSizes.data(),
                    scratch);
                break;
            }
            case VectorEncoding::Simple::LAZY:
                estimateSerializedSizeInt(vector->loaded_vector(), ranges, sizes, scratch);
                break;
            default:
                POLLUX_UNSUPPORTED("Unsupported vector encoding {}", vector->encoding());
        }
    }

    void estimateWrapperSerializedSize(
        const melon::Range<const IndexRange *> &ranges,
        vector_size_t **sizes,
        const BaseVector *wrapper,
        Scratch &scratch) {
        std::vector<IndexRange> newRanges;
        std::vector<vector_size_t *> newSizes;
        const BaseVector *wrapped = wrapper->wrapped_vector();
        for (int32_t i = 0; i < ranges.size(); ++i) {
            int32_t numNulls = 0;
            auto end = ranges[i].begin + ranges[i].size;
            for (int32_t offset = ranges[i].begin; offset < end; ++offset) {
                if (!wrapper->is_null_at(offset)) {
                    newRanges.push_back(IndexRange{wrapper->wrapped_index(offset), 1});
                    newSizes.push_back(sizes[i]);
                } else {
                    ++numNulls;
                }
            }
            *sizes[i] += bits::nbytes(numNulls);
        }
        estimateSerializedSizeInt(wrapped, newRanges, newSizes.data(), scratch);
    }

    void expandRepeatedRanges(
        const BaseVector *vector,
        const vector_size_t *rawOffsets,
        const vector_size_t *rawSizes,
        const melon::Range<const IndexRange *> &ranges,
        vector_size_t **sizes,
        std::vector<IndexRange> *childRanges,
        std::vector<vector_size_t *> *childSizes) {
        for (int32_t i = 0; i < ranges.size(); ++i) {
            int32_t begin = ranges[i].begin;
            int32_t end = begin + ranges[i].size;
            bool hasNull = false;
            for (int32_t offset = begin; offset < end; ++offset) {
                if (vector->is_null_at(offset)) {
                    hasNull = true;
                } else {
                    // Add the size of the length.
                    *sizes[i] += sizeof(int32_t);
                    childRanges->push_back(
                        IndexRange{rawOffsets[offset], rawSizes[offset]});
                    childSizes->push_back(sizes[i]);
                }
            }

            if (hasNull) {
                // Add the size of the null bit mask.
                *sizes[i] += bits::nbytes(ranges[i].size);
            }
        }
    }

    void estimateSerializedSizeInt(
        const BaseVector *vector,
        const melon::Range<const vector_size_t *> &rows,
        vector_size_t **sizes,
        Scratch &scratch) {
        const auto numRows = rows.size();
        if (vector->type()->isFixedWidth() && !vector->may_have_nulls_recursive()) {
            const auto elementSize = vector->type()->cppSizeInBytes();
            for (auto i = 0; i < numRows; ++i) {
                *sizes[i] += elementSize;
            }
            return;
        }
        switch (vector->encoding()) {
            case VectorEncoding::Simple::FLAT: {
                POLLUX_DYNAMIC_TYPE_DISPATCH_ALL(
                    estimateFlatSerializedSize,
                    vector->type_kind(),
                    vector,
                    rows,
                    sizes,
                    scratch);
                break;
            }
            case VectorEncoding::Simple::CONSTANT:
                POLLUX_DYNAMIC_TYPE_DISPATCH_ALL(
                    estimateFlattenedConstantSerializedSize,
                    vector->type_kind(),
                    vector,
                    rows,
                    sizes,
                    scratch);
                break;
            case VectorEncoding::Simple::DICTIONARY:
            case VectorEncoding::Simple::SEQUENCE:
                estimateWrapperSerializedSize(rows, sizes, vector, scratch);
                break;
            case VectorEncoding::Simple::BIASED:
                estimateBiasedSerializedSize(vector, rows, sizes, scratch);
                break;
            case VectorEncoding::Simple::ROW: {
                ScratchPtr<vector_size_t, 1> innerRowsHolder(scratch);
                ScratchPtr<vector_size_t *, 1> innerSizesHolder(scratch);
                ScratchPtr<uint64_t, 1> nullsHolder(scratch);
                auto *innerRows = rows.data();
                auto *innerSizes = sizes;
                const auto numRows = rows.size();
                int32_t numInner = numRows;
                if (vector->may_have_nulls()) {
                    auto nulls = nullsHolder.get(bits::nwords(numRows));
                    simd::gatherBits(vector->raw_nulls(), rows, nulls);
                    auto mutableInnerRows = innerRowsHolder.get(numRows);
                    numInner = simd::indicesOfSetBits(nulls, 0, numRows, mutableInnerRows);
                    innerSizes = innerSizesHolder.get(numInner);
                    for (auto i = 0; i < numInner; ++i) {
                        innerSizes[i] = sizes[mutableInnerRows[i]];
                    }
                    simd::transpose(
                        rows.data(),
                        melon::Range<const vector_size_t *>(mutableInnerRows, numInner),
                        mutableInnerRows);
                    innerRows = mutableInnerRows;
                }
                auto *row_vector = vector->as<RowVector>();
                auto &children = row_vector->children();
                for (auto &child: children) {
                    if (child) {
                        estimateSerializedSizeInt(
                            child.get(),
                            melon::Range(innerRows, numInner),
                            innerSizes,
                            scratch);
                    }
                }
                break;
            }
            case VectorEncoding::Simple::MAP: {
                auto *map_vector = vector->as_unchecked<MapVector>();
                ScratchPtr<IndexRange> rangeHolder(scratch);
                ScratchPtr<vector_size_t *> sizesHolder(scratch);
                const auto numRanges = rowsToRanges(
                    rows,
                    map_vector->raw_nulls(),
                    map_vector->rawOffsets(),
                    map_vector->rawSizes(),
                    sizes,
                    rangeHolder,
                    &sizesHolder,
                    nullptr,
                    scratch);
                if (numRanges == 0) {
                    return;
                }
                estimateSerializedSizeInt(
                    map_vector->mapKeys().get(),
                    melon::Range<const IndexRange *>(rangeHolder.get(), numRanges),
                    sizesHolder.get(),
                    scratch);
                estimateSerializedSizeInt(
                    map_vector->mapValues().get(),
                    melon::Range<const IndexRange *>(rangeHolder.get(), numRanges),
                    sizesHolder.get(),
                    scratch);
                break;
            }
            case VectorEncoding::Simple::ARRAY: {
                auto *array_vector = vector->as<ArrayVector>();
                ScratchPtr<IndexRange> rangeHolder(scratch);
                ScratchPtr<vector_size_t *> sizesHolder(scratch);
                const auto numRanges = rowsToRanges(
                    rows,
                    array_vector->raw_nulls(),
                    array_vector->rawOffsets(),
                    array_vector->rawSizes(),
                    sizes,
                    rangeHolder,
                    &sizesHolder,
                    nullptr,
                    scratch);
                if (numRanges == 0) {
                    return;
                }
                estimateSerializedSizeInt(
                    array_vector->elements().get(),
                    melon::Range<const IndexRange *>(rangeHolder.get(), numRanges),
                    sizesHolder.get(),
                    scratch);
                break;
            }
            case VectorEncoding::Simple::LAZY:
                estimateSerializedSizeInt(vector->loaded_vector(), rows, sizes, scratch);
                break;
            default:
                POLLUX_UNSUPPORTED("Unsupported vector encoding {}", vector->encoding());
        }
    }
} // namespace kumo::pollux::serializer::presto::detail
